import random
from ete4 import Tree

# Create a normal tree.
t = Tree('((H:0.3,I:0.1):0.5, A:1, (B:0.4,(C:0.5,(J:1.3, (F:1.2, D:0.1):0.5):0.5):0.5):0.5);')

print(t)
#  ╭─┬╴H
# ─┤ ╰╴I
#  ├╴A
#  ╰─┬╴B
#    ╰─┬╴C
#      ╰─┬╴J
#        ╰─┬╴F
#          ╰╴D

# Let's locate some nodes using the get common ancestor method
ancestor = t.common_ancestor('J', 'F', 'C')

# the search_nodes method (I take only the first match )
A = next(t.search_nodes(name='A'))

# and using the shorcut to finding nodes by name
C = t['C']
H = t['H']
I = t['I']

# Let's now add some custom features to our nodes. add_props can be
#  used to add many features at the same time.
C.add_props(vowel=False, confidence=1.0)
A.add_props(vowel=True, confidence=0.5)

ancestor.add_props(nodetype='internal')

# Or, using the oneliner notation
t['H'].add_props(vowel=False, confidence=0.2)

# But we can automatize this. (note that i will overwrite the previous
# values)
for leaf in t.traverse():
    if leaf.name and leaf.name in 'AEIOU':
        leaf.add_props(vowel=True, confidence=random.random())
    else:
        leaf.add_props(vowel=False, confidence=random.random())

# Now we use these information to analyze the tree.
print('This tree has', len(list(t.search_nodes(vowel=True))), 'vowel nodes')
print('Which are', [leaf.name for leaf in t.leaves() if leaf.props['vowel']])

# But features may refer to any kind of data, not only simple
# values. For example, we can calculate some values and store them
# within nodes.
#
# Let's detect leaf nodes under 'ancestor' with distance higher thatn
# 1. Note that I'm traversing a subtree which starts from 'ancestor'
matches = [leaf for leaf in ancestor.leaves() if leaf.dist > 1.0]

# And save this pre-computed information into the ancestor node
ancestor.add_prop('long_branch_nodes', matches)

# Prints the precomputed nodes
print('These are nodes under ancestor with long branches',
      [n.name for n in ancestor.props['long_branch_nodes']])

# We can also use the add_prop() method to dynamically add new features.
label = input('custom label:')
value = input('custom label value:')

ancestor.add_prop(label, value)
print('Ancestor has now the [', label, '] attribute with value [', value, ']')
